# -*- coding: utf-8 -*-


import numpy as np
import scipy.io as sio
import time
import HotAtomDist_3D
import os
import socket
os.system('chcp 65001')
Host_Name = socket.gethostname()
if Host_Name == 'VictorChengPC':
    Computer_Name = 'X:\\Work\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
elif Host_Name == 'calcenter2':
    Computer_Name = 'D:\\vc\\mail.ustc.edu.cn\\Rb87BEC - Files\\'

# from variables import (AtomDensityHot, version_count, save_address, m0, n0, l0, num_of_band)
Start_Time = time.time()
# mz_list = np.linspace(-0.5, 0.5, 101)
mz_list = np.array([-0.2])
for mz_index in range(len(mz_list)):
    mz = mz_list[mz_index]
    Data_Name = "%s\\Rb87Data\\3DSOCbandData\\190827\\" \
                "3x3x3;4x41x41x41;V_xz=1.770E_r,V_y=-3.540E_r,Oxy=1.022E_r,Ozy=-1.022E_r," \
                "T=1.000e-07K,n=3.000e+18m-3del_psi=1.000pi,mz=%.3fErno_states_v1" % (Computer_Name, mz)
    AtomTemperture = 100e-9
    AtomDensityHot = 3e18
    Data = np.load(Data_Name+'.npz')

    energies = Data['energies']
    ESigmaZ = Data['ESigmaZ']
    ku_grid = Data['ku_grid']
    ky_grid = Data['ky_grid']
    kv_grid = Data['kv_grid']
    Weyl_Points = Data['Weyl_Points']
    Number_of_WeylPoints = Data['Number_of_WeylPoints']
    parameters = Data['parameters']
    mz = Data['mz']

    (mu, ns, nsp) = HotAtomDist_3D.ChemicalPotential(energies, AtomTemperture, AtomDensityHot)
    if np.abs(np.sum(ns) - AtomDensityHot) / AtomDensityHot > 0.01:
        input('whut')
    Spin_Texture = np.sum(nsp * ESigmaZ, 0)
    Spin_Texture_yinted = np.sum(ESigmaZ * ns, axis=(0, 2)) / np.sum(ns, axis=(0, 2))
    Spin_Texture_y05 = np.squeeze(Spin_Texture[:, np.where(ky_grid == 0.5)[0], :])
    np.savez(Data_Name+'.npz',
             ku_grid=ku_grid, ky_grid=ky_grid, kv_grid=kv_grid,
             energies=energies, ESigmaZ=ESigmaZ, ns=ns, Spin_Texture=Spin_Texture,
             Spin_Texture_yinted=Spin_Texture_yinted, Spin_Texture_y05=Spin_Texture_y05,
             Weyl_Points=Weyl_Points, Number_of_WeylPoints=Number_of_WeylPoints,
             parameters=parameters, mz=mz)
    sio.savemat(Data_Name+'.mat',
                {'ku_grid': ku_grid, 'ky_grid': ky_grid, 'kv_grid': kv_grid,
                 'energies': energies, 'ESigmaZ': ESigmaZ, 'ns': ns, 'Spin_Texture': Spin_Texture,
                 'Spin_Texture_yinted': Spin_Texture_yinted, 'Spin_Texture_y05': Spin_Texture_y05,
                 'Weyl_Points': Weyl_Points, 'Number_of_WeylPoints': Number_of_WeylPoints,
                 'parameters': parameters, 'mz': mz})
End_Time = time.time()
print('Start: ', Start_Time)
print('End: ', End_Time)
print('Total: ', End_Time - Start_Time)
